package cn.tworice.watermark.util;

import org.springframework.mock.web.MockMultipartFile;
import org.springframework.web.multipart.MultipartFile;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;

public class ImageSteganography {
    /**
     * 写入内容到图片
     * @param image 图像文件
     * @param text 写入的文本信息
     * @return MultipartFile 图像文件
     * @throws IOException 读取图片文件时发生异常
     */
    public static MultipartFile hideTextInImage(MultipartFile image, String text) throws IOException {
        BufferedImage bufferedImage = hideTextInImage(ImageIO.read(image.getInputStream()), text);
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        ImageIO.write(bufferedImage, "png", baos);
        byte[] imageBytes = baos.toByteArray();
        return new MockMultipartFile(image.getName(), image.getOriginalFilename(), image.getContentType(), imageBytes);
    }

    /**
     * 读取信息从图片
     * @param image 图像文件
     * @return 文本信息
     * @throws IOException 从图片读取信息时发生异常
     */
    public static String extractTextFromImage(MultipartFile image) throws IOException {
        return extractTextFromImage(ImageIO.read(image.getInputStream()));
    }

    /**
     * 将文本信息隐藏到图片中
     * @param image 图像文件
     * @param text 写入的文本信息
     * @return 新的图像文件
     */
    public static BufferedImage hideTextInImage(BufferedImage image, String text) {
        int messageLength = text.length();
        int imageWidth = image.getWidth();
        int imageHeight = image.getHeight();
        int[] pixels = new int[imageWidth * imageHeight];
        image.getRGB(0, 0, imageWidth, imageHeight, pixels, 0, imageWidth);

        // 将文本长度编码到前32个像素中
        for (int i = 0; i < 32; i++) {
            int pixel = pixels[i] & 0xFFFFFFFE; // 清除最低位
            int bit = (messageLength >> (31 - i)) & 1; // 获取文本长度的一个位
            pixels[i] = pixel | bit; // 将这个位设置到像素的最低位
        }

        // 将文本信息编码到剩余的像素中
        for (int i = 0; i < messageLength; i++) {
            char c = text.charAt(i);
            for (int j = 0; j < 16; j++) {
                int pixelIndex = 32 + i * 16 + j;
                int pixel = pixels[pixelIndex] & 0xFFFFFFFE; // 清除最低位
                int bit = (c >> (15 - j)) & 1; // 获取字符的一个位
                pixels[pixelIndex] = pixel | bit; // 将这个位设置到像素的最低位
            }
        }

        BufferedImage newImage = new BufferedImage(imageWidth, imageHeight, BufferedImage.TYPE_INT_RGB);
        newImage.setRGB(0, 0, imageWidth, imageHeight, pixels, 0, imageWidth);
        return newImage;
    }

    /**
     * 将BufferedImage转换为MultipartFile
     * @param bufferedImage 要转换的BufferedImage对象
     * @param formatName 格式化后的图像格式，如png
     * @param fileName 文件名
     * @return MultipartFile对象
     * @throws IOException 如果创建MultipartFile对象时发生异常
     */
    public static MultipartFile bufferedImageToMultipartFile(BufferedImage bufferedImage, String formatName, String fileName) throws IOException {
        // 将BufferedImage写入到一个ByteArrayOutputStream中
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        ImageIO.write(bufferedImage, formatName, byteArrayOutputStream);
        byteArrayOutputStream.flush();

        // 使用MockMultipartFile创建一个MultipartFile对象
        MultipartFile multipartFile = new MockMultipartFile(fileName, fileName, "image/" + formatName, byteArrayOutputStream.toByteArray());

        byteArrayOutputStream.close();
        return multipartFile;
    }

    /**
     * 从图片中提取隐藏的文本信息
     * @param image 图像文件
     * @return 提取的隐藏文本信息
     */
    public static String extractTextFromImage(BufferedImage image) {
        int imageWidth = image.getWidth();
        int imageHeight = image.getHeight();
        int[] pixels = new int[imageWidth * imageHeight];
        image.getRGB(0, 0, imageWidth, imageHeight, pixels, 0, imageWidth);

        // 从前32个像素中解码文本长度
        int messageLength = 0;
        for (int i = 0; i < 32; i++) {
            messageLength = (messageLength << 1) | (pixels[i] & 1);
        }

        // 从剩余的像素中解码文本信息
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < messageLength; i++) {
            char c = 0;
            for (int j = 0; j < 16; j++) {
                int pixelIndex = 32 + i * 16 + j;
                c = (char) ((c << 1) | (pixels[pixelIndex] & 1));
            }
            sb.append(c);
        }

        return sb.toString();
    }

    public static void main(String[] args) throws IOException {
        BufferedImage originalImage = ImageIO.read(new File("C:\\Users\\17867\\Pictures\\Saved Pictures\\1.jpg"));
        String secretMessage = "Hello, world!";

        // 隐藏文本信息到图片中
        BufferedImage imageWithText = hideTextInImage(originalImage, secretMessage);
        ImageIO.write(imageWithText, "png", new File("C:\\Users\\17867\\Pictures\\Saved Pictures\\2.jpg"));

        // 从图片中提取文本信息
        BufferedImage imageToExtract = ImageIO.read(new File("C:\\Users\\17867\\Pictures\\Saved Pictures\\2.jpg"));
        String extractedMessage = extractTextFromImage(imageToExtract);
        System.out.println("Extracted Message: " + extractedMessage);
    }
}
